﻿Imports System.Management.Automation
Imports System.Management.Automation.Runspaces
Imports System.Collections.Generic
Imports System.Collections.ObjectModel
Imports System.Threading

Public Class PSExecute
    Implements IDisposable

    Public Enum PSExecuteStatus
        Initialized
        Executing
        Canceled
        Errored
        Completed
    End Enum

#Region "Variables"
    Public Event DataGenerated As EventHandler(Of PSExecuteStreamData)
    Private WithEvents myHost As New PSExecuteHost
    Private myPowerShell As PowerShell
    Private myStatus As PSExecuteStatus = PSExecuteStatus.Initialized
    Private forceVerbose As Boolean = False
    Private forceDebug As Boolean = False
    Private forceError As Boolean = False
    Private myScript As String
    Private myParameters As Hashtable
    Private myRunSpace As Runspace
    Private sharedRunSpace As Boolean = False 'shared runspace not implemented yet
    Private myCTS As New CancellationTokenSource
#End Region

#Region "Properties"
    Public ReadOnly Property Status As PSExecuteStatus
        Get
            Return Me.myStatus
        End Get
    End Property
    Public ReadOnly Property IsVerbose As Boolean
        Get
            Return Me.forceVerbose
        End Get
    End Property
    Public ReadOnly Property IsDebug As Boolean
        Get
            Return Me.forceDebug
        End Get
    End Property
    Public ReadOnly Property Script As String
        Get
            Return Me.myScript
        End Get
    End Property
    Public ReadOnly Property IsRunSpaceShared As Boolean
        Get
            Return Me.sharedRunSpace
        End Get
    End Property
    Public ReadOnly Property Parameters As Collections.Specialized.IOrderedDictionary
        Get
            Return Me.myParameters
        End Get
    End Property
#End Region

#Region "Constructors"
    Public Sub New(script As String, parameters As Hashtable, Optional stopOnError As Boolean = False, Optional showVerbose As Boolean = False, Optional showDebug As Boolean = False)
        Me.myScript = script
        Me.myParameters = parameters

        Me.forceVerbose = showVerbose
        Me.forceDebug = showDebug
        Me.forceError = stopOnError

        Me.myRunSpace = RunspaceFactory.CreateRunspace(Me.myHost)
        Me.myRunSpace.Open()
    End Sub

    Public Sub New(script As String, Optional stopOnError As Boolean = False, Optional showVerbose As Boolean = False, Optional showDebug As Boolean = False)
        Me.New(script, Nothing, stopOnError, showVerbose, showDebug)
    End Sub
#End Region

    Public Sub Execute()
        If Not Me.myStatus = PSExecuteStatus.Initialized Then Throw New InvalidOperationException("Can only call 'Execute' when object is initialized.")
        Try
            Me.myStatus = PSExecuteStatus.Executing
            Me.myPowerShell = PowerShell.Create()
            Me.myPowerShell.Runspace = Me.myRunSpace

            If Me.sharedRunSpace = False Then
                If Me.forceDebug Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("DebugPreference", "Continue")
                If Me.forceVerbose Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("VerbosePreference", "Continue")
                If Me.forceError Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("ErrorActionPreference", "Stop")
            End If

            Me.myPowerShell.AddScript(Me.myScript)
            If Me.myParameters IsNot Nothing AndAlso Me.myParameters.Count > 0 Then Me.myPowerShell.AddParameters(Me.myParameters)
            Me.myPowerShell.AddCommand("Out-Default")
            Me.myPowerShell.Commands.Commands(0).MergeMyResults(PipelineResultTypes.Error, PipelineResultTypes.Output)

            Me.myPowerShell.Invoke()
            Me.myStatus = PSExecuteStatus.Completed
        Catch rte As RuntimeException
            Me.myStatus = PSExecuteStatus.Errored
            RaiseEvent DataGenerated(Me, New PSExecuteStreamData(rte.ToString, PSExecuteStreamDataSource.Error))
        Finally
            Me.myPowerShell.Dispose()
            Me.myPowerShell = Nothing
        End Try
    End Sub

    Public Async Function ExecuteAsync(Optional BubbleCancellationError As Boolean = False) As Task
        If Not Me.myStatus = PSExecuteStatus.Initialized Then Throw New InvalidOperationException("Can only call 'Execute' when object is initialized.")
        Dim canceltoken As CancellationToken = Me.myCTS.Token
        Await Task.Run(Sub()
                           Try
                               Me.myStatus = PSExecuteStatus.Executing
                               Me.myPowerShell = PowerShell.Create()
                               Me.myPowerShell.Runspace = Me.myRunSpace

                               If Me.sharedRunSpace = False Then
                                   If Me.forceDebug Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("DebugPreference", "Continue")
                                   If Me.forceVerbose Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("VerbosePreference", "Continue")
                                   If Me.forceError Then Me.myPowerShell.Runspace.SessionStateProxy.SetVariable("ErrorActionPreference", "Stop")
                               End If

                               Me.myPowerShell.AddScript(Me.myScript)
                               If Me.myParameters IsNot Nothing AndAlso Me.myParameters.Count > 0 Then Me.myPowerShell.AddParameters(Me.myParameters)
                               Me.myPowerShell.AddCommand("Out-Default")
                               Me.myPowerShell.Commands.Commands(0).MergeMyResults(PipelineResultTypes.Error, PipelineResultTypes.Output)

                               Dim ah As IAsyncResult = Me.myPowerShell.BeginInvoke()

                               If WaitHandle.WaitAny({ah.AsyncWaitHandle, canceltoken.WaitHandle}) = 1 Then
                                   If canceltoken.IsCancellationRequested Then
                                       canceltoken.ThrowIfCancellationRequested()
                                   End If
                               End If

                               Me.myPowerShell.EndInvoke(ah)
                               Me.myStatus = PSExecuteStatus.Completed
                           Catch rte As RuntimeException
                               Me.myStatus = PSExecuteStatus.Errored
                               RaiseEvent DataGenerated(Me, New PSExecuteStreamData(rte.ToString, PSExecuteStreamDataSource.Error))
                           Catch oce As OperationCanceledException
                               Me.myStatus = PSExecuteStatus.Canceled
                               Me.myPowerShell.Stop()
                               If BubbleCancellationError Then
                                   Throw oce
                               Else
                                   RaiseEvent DataGenerated(Me, New PSExecuteStreamData("Script execution canceled.", PSExecuteStreamDataSource.Error))
                               End If
                           Finally
                               Me.myPowerShell.Dispose()
                               Me.myPowerShell = Nothing
                           End Try
                       End Sub, canceltoken)
    End Function

    Public Sub CancelExecute(Optional wait As Boolean = False)
        If Not Me.myStatus = PSExecuteStatus.Executing Then Throw New InvalidOperationException("Cannot call 'CancelExecute' when an execution is not currently in progress.")
        Me.myCTS.Cancel()
        Me.myCTS.Token.WaitHandle.WaitOne(100)
    End Sub

#Region "IDisposable Support"
    Private disposedValue As Boolean ' To detect redundant calls

    ' IDisposable
    Protected Overridable Sub Dispose(disposing As Boolean)
        If Not Me.disposedValue Then
            If disposing Then
                ' TODO: dispose managed state (managed objects).
                If Me.myStatus = PSExecuteStatus.Executing Then Me.CancelExecute(True)
                If Me.myPowerShell IsNot Nothing Then
                    Me.myPowerShell.Stop()
                    Me.myPowerShell.Dispose()
                    Me.myPowerShell = Nothing
                End If
                If Not Me.sharedRunSpace Then
                    Me.myRunSpace.Close()
                    Me.myRunSpace.Dispose()
                    Me.myRunSpace = Nothing
                End If
                Me.myCTS.Dispose()
            End If

            ' TODO: free unmanaged resources (unmanaged objects) and override Finalize() below.
            ' TODO: set large fields to null.
        End If
        Me.disposedValue = True
    End Sub

    ' TODO: override Finalize() only if Dispose(ByVal disposing As Boolean) above has code to free unmanaged resources.
    'Protected Overrides Sub Finalize()
    '    ' Do not change this code.  Put cleanup code in Dispose(ByVal disposing As Boolean) above.
    '    Dispose(False)
    '    MyBase.Finalize()
    'End Sub

    ' This code added by Visual Basic to correctly implement the disposable pattern.
    Public Sub Dispose() Implements IDisposable.Dispose
        ' Do not change this code.  Put cleanup code in Dispose(disposing As Boolean) above.
        Dispose(True)
        GC.SuppressFinalize(Me)
    End Sub
#End Region

    Private Sub myHost_DataGenerated(sender As Object, e As PSExecuteStreamData) Handles myHost.DataGenerated
        RaiseEvent DataGenerated(Me, e)
    End Sub

    Public Shared Function Caculate(Of TInput, TResult)(script As String, input As TInput, Optional millisecondLimit As Integer = 250) As TResult
        Dim result As TResult = Nothing
        Dim ps As PowerShell = PowerShell.Create()
        ps.Runspace = RunspaceFactory.CreateRunspace()
        ps.Runspace.Open()

        Try
            script = script.Trim()
            If Not script.StartsWith("param(") Then script = "param($_) " + script
            ps.AddScript(script).AddArgument(input)

            Dim ar As IAsyncResult = ps.BeginInvoke()
            If ar.AsyncWaitHandle.WaitOne(millisecondLimit, True) Then

                Using tempPSD As PSDataCollection(Of PSObject) = ps.EndInvoke(ar)
                    If tempPSD.Count > 0 Then
                        result = CType(tempPSD.Item(0).BaseObject, TResult)
                    End If
                End Using
            Else
                ps.Stop()
                Throw New TimeoutException(String.Format("Timeout ({0} milliseconds) elapsed for PSEngine.Calculate."))
            End If
        Catch ex As Exception
            Throw ex
        Finally
            ps.Runspace.Close()
            ps.Dispose()
            ps = Nothing
        End Try

        Return result
    End Function
End Class
